import pygame, sys
from pygame.locals import *
import math

FPS = 30 # frames per second setting

WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED   = (255, 0, 0)
GREEN = (0, 255, 0)
GRAY  = pygame.Color('Gray')

orgx = 100
orgy = 100

sin_x = orgx
sin_y = orgy+200

start = 0
''' 公式: y = A * sin(w*x + f) + b '''

''' 周期 '''
w = 1.0
T = 360 / w
A = 100

''' 计算周期函数里每个x对应的sin值 '''
def MakeSinArr(w, T):
    sinArr = {}
    for degree in range(0, int(T)+1):
        ''' 角度转换为弧度值 '''
        rad = math.radians(degree)
        ''' 计算 sin(rad) 的值'''
        y = math.sin(rad * w)
        ''' 添加到字典里 '''
        sinArr[degree] = y
    return sinArr

''' 画sin 曲线 '''
def DawSinCurve(sinArr, offset, color):
    global sin_y
    for x in range(orgx, 731, 1):
        y = sin_y - A * sinArr[int((x+offset-orgx)%T)]
        pygame.draw.circle(DISPLAYSURF, color, (x, y), 1) 

'''  初始化配置 '''
def initEvn():
    global start, sin_x, sin_y
    start = 0
    sin_y = orgy + 200
    sin_x = orgx
    
def drawBg():
    blockWidth = 90
    blockHeight = 100
    gridWidth = 90 * 7
    gridHeight = 100 * 4
    
    ''' 填充背景 白色 '''
    DISPLAYSURF.fill(pygame.Color("WHITE"))

    ''' 画竖线 '''
    for x in range(orgx, gridWidth + orgx + 1, 10):
        if not (x-orgx) % blockWidth:
            pygame.draw.line(DISPLAYSURF, BLACK, (x, orgy), (x, orgy + gridHeight), 1)
        else:
            pygame.draw.line(DISPLAYSURF, GRAY, (x, orgy), (x, orgy + gridHeight), 1)

    ''' 画横线 '''
    for y in range(orgy, gridHeight + orgy + 1, 10):
        if not (y-orgy) % blockHeight:
            pygame.draw.line(DISPLAYSURF, BLACK, (orgx, y), (orgx + gridWidth, y), 1)
        else:
            pygame.draw.line(DISPLAYSURF, GRAY, (orgx, y), (orgx + gridWidth, y), 1)

    ''' 画横纵轴 '''
    xStart = (orgx-20, orgy+200)
    xEnd = (orgx + gridWidth +20, orgy+200)

    yStart = (orgx, orgy-20)
    yEnd = (orgx, orgy + gridHeight +20)
    
    pygame.draw.line(DISPLAYSURF, BLACK, xStart, xEnd, 2)
    pygame.draw.line(DISPLAYSURF, BLACK, yStart, yEnd, 2)

    ''' 画横纵轴箭头 '''
    pygame.draw.line(DISPLAYSURF, BLACK, (orgx + gridWidth + 20 - 5, orgy+200-5), xEnd, 2)
    pygame.draw.line(DISPLAYSURF, BLACK, (orgx + gridWidth + 20 - 5 , orgy+200+5), xEnd, 2)

    pygame.draw.line(DISPLAYSURF, BLACK, (orgx-5, orgy-15), yStart, 2)
    pygame.draw.line(DISPLAYSURF, BLACK, (orgx+5, orgy-15), yStart, 2)

    ''' 画坐标单位值 '''
    for y in range(orgy, orgy+gridHeight+1, 100):
        yUnitSurf = BASICFONT.render('%3d'%(300-y), True, BLACK)
        yUnitRect = yUnitSurf.get_rect()
        yUnitRect.topleft = (orgx - 44, y -10)
        DISPLAYSURF.blit(yUnitSurf, yUnitRect)
    for x in range(orgx+90, orgx+gridWidth+1, 90):
        xUnitSurf = BASICFONT.render('%3d'%(x-100), True, BLACK)
        xUnitRect = xUnitSurf.get_rect()
        xUnitRect.topleft = (x - 22, orgy -20)
        DISPLAYSURF.blit(xUnitSurf, xUnitRect)
    
'''  画函数信息， 振幅A 周期T 角速度w '''
def drawInfo(A,w,T):
    infoSurf = BASICFONT.render(' A: %6d  w: %5.2f  T: %8.2f ' % (abs(A), w, T), True, BLACK)
    infoRect = infoSurf.get_rect()
    infoRect.topleft = (250, 20)

    sinSurf = BASICFONT.render('sin curve', True, RED)
    sinRect = sinSurf.get_rect()
    sinRect.topleft = (50, 20)

    cosSurf = BASICFONT.render('cos curve', True, GREEN)
    cosRect = cosSurf.get_rect()
    cosRect.topleft = (50, 40)
    
    DISPLAYSURF.blit(sinSurf, sinRect)
    DISPLAYSURF.blit(cosSurf, cosRect)
    DISPLAYSURF.blit(infoSurf, infoRect)
    
def terminate():
    pygame.quit()
    sys.exit()

def main():
    global FPSCLOCK, DISPLAYSURF, BASICFONT, A, T, w, start
    pygame.init()
    
    fpsClock = pygame.time.Clock()

    # set up the window
    DISPLAYSURF = pygame.display.set_mode((800, 600), 0, 32)
    pygame.display.set_caption('模拟Sin Cos曲线')

    BASICFONT = pygame.font.Font('freesansbold.ttf', 20)
    ''' 画背景网格 '''
    drawBg()

    ''' 读取周期函数x与y的 映射 '''
    sinArr = MakeSinArr(w, T)
    ''' 主循环 '''
    while True: # the main game loop
    
        ''' 画网格背景 '''
        drawBg()
    
        ''' 画函数信息 '''
        drawInfo(A,w,T)
    
        ''' 画正弦函数曲线 '''
        DawSinCurve(sinArr, start, pygame.Color('Red'))

        ''' 画余弦函数曲线 '''
        DawSinCurve(sinArr, start+T/4.0, pygame.Color('Green'))

        ''' 对称正弦 '''
        DawSinCurve(sinArr, start + T/2.0, pygame.Color('Blue'))
    
        ''' 对称余弦 '''
        DawSinCurve(sinArr, start - T/4.0, pygame.Color('Yellow'))
    
        '''
            动起来
            走1个周期, 结束之后可以按空格再次触发
            也可以将T 替换成 T * 1000 走1000个周期
        '''
        if start < (T*1000):
           start += 2
        
        for event in pygame.event.get():
            ''' 键盘事件 '''
            if (event.type == KEYDOWN):
                ''' 空格键按下 且 不在运动中 '''
                if event.key == K_SPACE and start >= T:
                    initEvn()
                    ''' 方向键 上 下 或者 w s 按下： 调整振幅 A '''
                elif event.key == K_UP or event.key == K_w:
                    A += 10
                elif event.key == K_DOWN or event.key == K_s:
                    A -= 10
                    ''' 方向键 左 右 或者 a d 按下： 调整周期 '''
                elif event.key == K_LEFT  or event.key == K_a:
                    w += 0.1
                    T = 360/w
                    sinArr = MakeSinArr(w, T)
                elif event.key == K_RIGHT  or event.key == K_d:
                    w -= 0.1
                    T = 360/w
                    sinArr = MakeSinArr(w, T)
                elif event.key == K_ESCAPE:
                    terminate()
            elif event.type == QUIT:
                terminate()

        pygame.display.update()
        fpsClock.tick(FPS)

if __name__ == '__main__':
    main()
